from tqdm import tqdm
import torch
import torchvision as tv
from torch.utils.data import DataLoader
import torch.nn as nn
from Config import Config
from Model.Generation import Generation
from Model.Discrimination import Discrimination

opt = Config()

@torch.no_grad()
def generate(**kwargs):
    # 用训练好的模型来生成图片

    for k_, v_ in kwargs.items():
        setattr(opt, k_, v_)

    device = torch.device("cuda") if opt.gpu else torch.device("cpu")

    # 加载训练好的权重数据
    netg, netd = Generation(opt).eval(), Discrimination(opt).eval()
    #  两个参数返回第一个
    map_location = lambda storage, loc: storage

    # opt.netd_path等参数有待修改
    netd.load_state_dict(torch.load('save_img/netd_399.pth', map_location=map_location), False)
    netg.load_state_dict(torch.load('save_img/netg_399.pth', map_location=map_location), False)
    netd.to(device)
    netg.to(device)

    # 生成训练好的图片
    # 初始化512组噪声，选其中好的拿来保存输出。
    noise = torch.randn(opt.gen_search_num, opt.nz, 1, 1).normal_(opt.gen_mean, opt.gen_std).to(device)

    fake_image = netg(noise)
    score = netd(fake_image).detach()

    # 挑选出合适的图片
    # 取出得分最高的图片
    indexs = score.topk(opt.gen_num)[1]

    result = []

    for ii in indexs:
        result.append(fake_image.data[ii])

    # 以opt.gen_img为文件名保存生成图片
    tv.utils.save_image(torch.stack(result), opt.gen_img, normalize=True, range=(-1, 1))